/*
 * CMediaWriter.cpp
 *
 *  Created on: 2017年3月5日
 *      Author: chuanjiang.zh
 */

#include "CMediaWriter.h"
#include "CAudioTranscoder.h"
#include "MFormatUtil.h"
#include "CAudioTranscoder.h"
#include "CLog.h"
#include "TStringCast.h"
#include "TFileUtil.h"

namespace av
{

CMediaWriter::CMediaWriter():
		m_fmt(),
		m_fileName(),
		m_changeFile(),
		m_fmtCtx(),
		m_ioCtx(),
		m_videoIndex(-1),
		m_audioIndex(-1),
		m_videoStartTime(AV_NOPTS_VALUE),
		m_videoCurrentTime(),
		m_audioStartTime(AV_NOPTS_VALUE),
		m_audioCurrentTime(),
		m_transcoder(new CAudioTranscoder()),
		m_doAudioTranscode(false),
		m_audioCodecId(AV_CODEC_ID_NONE),
		m_byteCount(),
		m_listener()
{
	m_audioBuffer.ensure(1024 * 100);
}

CMediaWriter::~CMediaWriter()
{
	close();
}

bool CMediaWriter::open(const char* filename, const MFormat& fmt)
{
	if (!filename || strlen(filename) == 0)
	{
		return false;
	}
	
	//std::string audioConfig = comn::StringCast::toHexString(fmt.config, fmt.configSize);
	//CLog::debug("audioConfig: %s\n", audioConfig.c_str());
					
	m_fileName = filename;
	MFormatUtil::copyTo(&fmt, m_fmt);

	if (!m_fmt.isValid())
	{
		return false;
	}

	if (m_audioCodecId != AV_CODEC_ID_NONE)
	{
		m_doAudioTranscode = (m_audioCodecId != fmt.audioCodec);
	}

	//CLog::debug("CMediaWriter::open %s. video codec:%d, audio codec:%d, channels:%d, rate:%d. target codec:%d. trans:%d\n", 
	//	filename, fmt.codec,
	//	fmt.audioCodec, fmt.channels, fmt.sampleRate, m_audioCodecId, m_doAudioTranscode);

	return true;
}

void CMediaWriter::close()
{
	if (isFileOpen())
	{
		//CLog::debug("CMediaWriter::close %s\n", m_fileName.c_str());
	}
	
	closeFile();
}

bool CMediaWriter::isOpen()
{
	return isFileOpen();
}

bool CMediaWriter::getFormat(MFormat& fmt)
{
	MFormatUtil::copyTo(&fmt, m_fmt);

	return true;
}

bool CMediaWriter::writeVideo(uint8_t* data, int size, int64_t pts, int flags)
{
	//comn::FileUtil::write(data, size, "dump_2.h264", true);

	//CLog::debug("video packet. size:%d, flags:%d, dts:%lld\n",  size, flags, pts);

	if (m_changeFile)
	{
		if ((flags & AV_PKT_FLAG_KEY) != 0)
		{
			closeFile();

			m_changeFile = false;

			if (openFile())
			{
				fireStatus(MEDIA_WRITER_OPENED);
			}
			else
			{
				fireStatus(MEDIA_WRITER_IO_ERROR);
			}
		}
	}

	if (!isFileOpen() && !m_fileName.empty())
	{
		if (openFile())
		{
			fireStatus(MEDIA_WRITER_OPENED);
		}
		else
		{
			fireStatus(MEDIA_WRITER_IO_ERROR);
		}
	}

	if (m_videoIndex < 0 || !isFileOpen())
	{
		CLog::warning("file is not ready. size:%d, flags:%d\n", size, flags);
		return false;
	}

	AVPacket pkt;
	av_init_packet(&pkt);
	pkt.data = data;
	pkt.size = size;
	pkt.pts = pts;
	pkt.dts = pts;
	pkt.flags = flags;
	pkt.stream_index = m_videoIndex;

	writeVideo(pkt);

	return true;
}

bool CMediaWriter::writeAudio(uint8_t* data, int size, int64_t pts, int flags)
{
	if (m_audioIndex < 0 || !isFileOpen())
	{
		if (!m_fmt.hasVideo() && m_fmt.hasAudio())
		{
			if (openFile())
			{
				fireStatus(MEDIA_WRITER_OPENED);
			}
			else
			{
				fireStatus(MEDIA_WRITER_IO_ERROR);
				return false;
			}
		}
		else
		{
			return false;
		}
	}
	
	AVPacket pkt;
	av_init_packet(&pkt);
	pkt.data = data;
	pkt.size = size;
	pkt.pts = pts;
	pkt.dts = pts;
	pkt.flags = flags;
	pkt.stream_index = m_audioIndex;

	writeAudio(pkt);

	return true;
}

int CMediaWriter::getDuration()
{
	int64_t duration = m_videoCurrentTime - m_videoStartTime;

	return (int)(duration / AV_TIME_BASE);
}

int64_t CMediaWriter::getByteLength()
{
	return m_byteCount;
}

bool CMediaWriter::setFile(const char* filename)
{
    if (filename == NULL)
    {
        m_fileName.clear();
    }
    else
    {
        m_fileName = filename;
    }

    {
        comn::AutoCritSec lock(m_cs);
        if (isFileOpen())
        {
            m_changeFile = true;
        }
    }

    return true;
}

bool CMediaWriter::setOutputAudioCodec(int codecId)
{
	//CLog::debug("setOutputAudioCodec. codec:%d\n", codecId);
	
	bool result = false;
	if (codecId == AV_CODEC_ID_AAC)
	{
		m_audioCodecId = AV_CODEC_ID_AAC;
		result = true;
	}
	else
	{
		m_audioCodecId = AV_CODEC_ID_NONE;
	}

	return result;
}

void CMediaWriter::setListener(MediaWriterListener* listener)
{
	m_listener = listener;
}

bool CMediaWriter::openIOContext(const char* filename)
{
	avio_open(&m_ioCtx, filename, AVIO_FLAG_WRITE);
	return (m_ioCtx != NULL);
}

void CMediaWriter::closeIOContext()
{
    if (m_ioCtx)
    {
        avio_close(m_ioCtx);
        m_ioCtx = NULL;
    }
}


bool CMediaWriter::openFile()
{
	comn::AutoCritSec lock(m_cs);
	
    int rc = avformat_alloc_output_context2(&m_fmtCtx, NULL, NULL, m_fileName.c_str());
    if (!m_fmtCtx)
    {
		CLog::warning("can not alloc avformat.\n");
        return false;
    }

    if (m_fmt.hasVideo())
    {
        AVCodecID codecID = (AVCodecID)m_fmt.m_codec;
        AVCodec* codec = avcodec_find_decoder(codecID);
        if (codec != NULL)
        {
			//CLog::debug("video codec: %s\n", codec->name);
			
            AVStream *out_stream = avformat_new_stream(m_fmtCtx, NULL);
            AVCodecContext * codecCtx = out_stream->codec;

            avcodec_get_context_defaults3(codecCtx, codec);

            if (m_fmt.m_width <= 0)
            {
                codecCtx->width = 1920;
                codecCtx->height = 1080;
            }
            else
            {
                codecCtx->width = m_fmt.m_width;
                codecCtx->height = m_fmt.m_height;
                codecCtx->profile = m_fmt.m_profile;
            }

            codecCtx->time_base = av_make_q(1, 90000);

            out_stream->time_base = codecCtx->time_base;

            // extra_data
            if (m_fmt.m_videoProp.size() > 0)
            {
                codecCtx->extradata = (uint8_t*)av_memdup(m_fmt.m_videoProp.c_str(), m_fmt.m_videoProp.size());
                codecCtx->extradata_size = m_fmt.m_videoProp.size();
            }

            if (m_fmtCtx->oformat->flags & AVFMT_GLOBALHEADER)
            {
                out_stream->codec->flags |= CODEC_FLAG_GLOBAL_HEADER;
            }

            m_videoIndex = 0;
        }
		else
		{
			CLog::warning("no such video codec %d\n", codecID);
		}
    }

    if (m_fmt.hasAudio())
    {
        AVCodecID codecID = (AVCodecID)m_fmt.m_audioCodec;

        if (isG711(m_fmt) && isMp4(m_fmtCtx))
        {
        	m_audioCodecId = AV_CODEC_ID_AAC;
        	m_doAudioTranscode = true;
        }

        if (m_doAudioTranscode)
        {
			//CLog::debug("open transcoder. codec:%d, channels:%d, rate:%d\n", 
			//	m_audioCodecId, m_fmt.m_channels, m_fmt.m_sampleRate);
			m_transcoder->open(m_fmt, m_audioCodecId, m_fmt.m_channels, m_fmt.m_sampleRate);
        }

        AVCodec* codec = avcodec_find_decoder(codecID);
        if (codec != NULL)
        {
			//CLog::debug("audio codec: %s\n", codec->name);
			
            AVStream *out_stream = avformat_new_stream(m_fmtCtx, codec);
            AVCodecContext * codecCtx = out_stream->codec;

            avcodec_get_context_defaults3(codecCtx, codec);

            if (m_transcoder && m_transcoder->isOpen())
            {
                avcodec_copy_context(codecCtx, m_transcoder->getCodecContext());

                out_stream->time_base = codecCtx->time_base;
                if (m_fmtCtx->oformat->flags & AVFMT_GLOBALHEADER)
                {
                    out_stream->codec->flags |= CODEC_FLAG_GLOBAL_HEADER;
                }

            }
            else
            {
                codecCtx->channels = m_fmt.m_channels;
                codecCtx->channel_layout = av_get_default_channel_layout(m_fmt.m_channels);
                codecCtx->sample_rate = m_fmt.m_sampleRate;
                codecCtx->profile = m_fmt.m_audioProfile;
                codecCtx->time_base.num = 1;
                codecCtx->time_base.den = codecCtx->sample_rate;

                out_stream->time_base = codecCtx->time_base;

                if (m_fmt.m_audioConfig.size() > 0)
                {
                    codecCtx->extradata = (uint8_t*)av_memdup(m_fmt.m_audioConfig.c_str(), m_fmt.m_audioConfig.size());
                    codecCtx->extradata_size = m_fmt.m_audioConfig.size();
					
					//std::string extra = comn::StringCast::toHexString(codecCtx->extradata, //codecCtx->extradata_size);
					//CLog::debug("extra: %s\n", extra.c_str());
                }

                if (m_fmtCtx->oformat->flags & AVFMT_GLOBALHEADER)
                {
                    out_stream->codec->flags |= CODEC_FLAG_GLOBAL_HEADER;
                }
            }

            m_audioIndex = m_videoIndex + 1;
        }
    }

    rc = avio_open(&m_fmtCtx->pb, m_fileName.c_str(), AVIO_FLAG_WRITE);
    if (m_fmtCtx->pb == NULL)
    {
		CLog::error("failed to open file %s\n", m_fileName.c_str());

        avformat_free_context(m_fmtCtx);
        m_fmtCtx = NULL;
        m_videoIndex = -1;
        m_audioIndex = -1;
        return false;
    }

    rc = avformat_write_header(m_fmtCtx, NULL);
	if (rc != 0)
	{
		CLog::error("failed to write header %s\n", m_fileName.c_str());

		avformat_free_context(m_fmtCtx);
		m_fmtCtx = NULL;
		m_videoIndex = -1;
		m_audioIndex = -1;
		return false;
	}

	return true;
}

void CMediaWriter::closeFile()
{
	if (isFileOpen())
	{
		fireStatus(MEDIA_WRITER_CLOSED);
	}

	comn::AutoCritSec lock(m_cs);
	
	if (m_fmtCtx)
	{
		if (m_fmtCtx->pb)
		{
			av_write_trailer(m_fmtCtx);
			avio_close(m_fmtCtx->pb);
		}

		for (unsigned int i = 0; i < m_fmtCtx->nb_streams; i ++)
		{
			AVStream* stream = m_fmtCtx->streams[i];
			if (stream && stream->codec)
			{
				avcodec_close(stream->codec);
			}
		}

        avformat_free_context(m_fmtCtx);
        m_fmtCtx = NULL;
    }

	if (m_transcoder)
	{
		m_transcoder->close();
	}

    m_changeFile = false;

    m_videoIndex = -1;
    m_audioIndex = -1;
    m_videoStartTime = AV_NOPTS_VALUE;
    m_audioStartTime = AV_NOPTS_VALUE;
    m_videoCurrentTime = 0;
    m_audioCurrentTime = 0;

    m_byteCount = 0;
}


bool CMediaWriter::isFileOpen()
{
    return (m_fmtCtx != NULL);
}


bool CMediaWriter::isG711(const MediaFormat& fmt)
{
    return (fmt.m_audioCodec == AV_CODEC_ID_PCM_ALAW) || (fmt.m_audioCodec == AV_CODEC_ID_PCM_MULAW);
}

bool CMediaWriter::isMp4(AVFormatContext* fmtCtx)
{
    int order = strcmp(fmtCtx->oformat->name, "mp4");
    return (order == 0);
}


void CMediaWriter::writeVideo(AVPacket& inPkt)
{
    if (m_videoIndex < 0)
    {
		CLog::warning("file is not ready. video index < 0. size:%d, flags:%d\n", inPkt.size, inPkt.flags);
        return;
    }
	
	comn::AutoCritSec lock(m_cs);

    m_byteCount += inPkt.size;

    if (m_videoStartTime == AV_NOPTS_VALUE)
    {
        m_videoStartTime = inPkt.dts;
        m_videoCurrentTime = m_videoStartTime;
    }

    int64_t dts = inPkt.dts - m_videoStartTime;

    int64_t duration = inPkt.dts - m_videoCurrentTime;
    if (duration < 0)
    {
    	duration = 0;
    }

    m_videoCurrentTime = inPkt.dts;

    AVPacket pkt;
    av_init_packet(&pkt);

    av_packet_ref(&pkt, &inPkt);

    pkt.stream_index = m_videoIndex;
    pkt.pts = dts;
    pkt.dts = dts;
    pkt.duration = duration;

    AVRational src = av_make_q(1, AV_TIME_BASE);
    AVRational dst = m_fmtCtx->streams[m_videoIndex]->time_base;
    av_packet_rescale_ts(&pkt, src, dst);

    //CLog::debug("video packet. fmtCtx:%p, size:%d, duration:%d, dts:%lld\n", 
	//	m_fmtCtx, pkt.size, (int)pkt.duration, pkt.dts);

    av_interleaved_write_frame(m_fmtCtx, &pkt);

    av_packet_unref(&pkt);
}

void CMediaWriter::writeAudio(AVPacket& inPkt)
{
	//CLog::debug("CMediaWriter::writeAudio. size:%d, dts:%lld\n",  inPkt.size, inPkt.dts);

    if (m_doAudioTranscode)
    {
        AVPacket outPkt;
        av_init_packet(&outPkt);
        outPkt.data = m_audioBuffer.data();
        outPkt.size = m_audioBuffer.capacity();

        if (m_transcoder->transcode(&inPkt, &outPkt))
        {
            outPkt.stream_index = m_audioIndex;
            doWriteAudio(&outPkt, m_fmt.m_sampleRate);

            while (true)
            {
                av_init_packet(&outPkt);
                outPkt.data = m_audioBuffer.data();
                outPkt.size = m_audioBuffer.capacity();

                if (!m_transcoder->transcode(NULL, &outPkt))
                {
                    break;
                }

                outPkt.stream_index = m_audioIndex;

                doWriteAudio(&outPkt, m_fmt.m_sampleRate);
            }
        }
    }
    else
    {
        if (m_audioStartTime == AV_NOPTS_VALUE)
        {
            m_audioStartTime = inPkt.pts;
        }

        int64_t pts = inPkt.pts - m_audioStartTime;

        AVPacket pkt;
        av_init_packet(&pkt);

        av_packet_ref(&pkt, &inPkt);
        pkt.stream_index = m_audioIndex;
        pkt.pts = pts;
        pkt.dts = pts;

        doWriteAudio(&pkt, AV_TIME_BASE);

        av_packet_unref(&pkt);
    }
}

void CMediaWriter::doWriteAudio(AVPacket* pkt, int rate)
{
	//CLog::debug("Audio packet. size:%d, dts:%lld\n",  pkt->size, pkt->dts);

	comn::AutoCritSec lock(m_cs);
	
	m_byteCount += pkt->size;

	if(m_fmtCtx)
	{
		AVRational src = av_make_q(1, rate);
		AVRational dst = m_fmtCtx->streams[m_audioIndex]->time_base;
		av_packet_rescale_ts(pkt, src, dst);

		av_interleaved_write_frame(m_fmtCtx, pkt);
	}
	
	//CLog::debug("Audio packet end.\n");
}

void CMediaWriter::fireStatus(int status)
{
	if (m_listener)
	{
		m_listener->onWriteStatus(this, status);
	}
}

bool CMediaWriter::setMeta(const char* key, const char* value, int size)
{
	return true;
}


} /* namespace av */
